import logging
import struct
import importlib
from pack_to_json.protocol.frame_pb2 import Frame as Frame_v1, FrameShuffle
from pack_to_json.protocol.frame.v2_pb2 import Frame as Frame_v2

END = -1
CONTINUE = 1
SUCCESS = 2

TOPIC_PROTOS = {
    # perception
    "adas_function": ["adas_function_pb2", "AdasFunctionResult"],
    "adas_output_data": ["adas_output_info2_pb2", "ADASOutputFrame"],
    "calib_mono_extrin": ["calib_mono_extrin_pb2", "CalibMonoExtrinProto"],
    "diag_info": ["diag_info_pb2", "DiagInfo"],
    "weather_type": ["work_condition_pb2", "WorkCondition"],
    "scene_type": ["work_condition_pb2", "WorkCondition"],
    "light_type": ["work_condition_pb2", "WorkCondition"],
    "time_type": ["work_condition_pb2", "WorkCondition"],
    "ihbc": ["ihbc_output_pb2", "IHBCResultProto"],
    "image_quality": ["image_fail_pb2", "ImageFail"],
    "image": ["image_pb2", "Image"],
    "isp_param": ["isp_param_pb2", "ISPParamProto"],
    "log": ["log_pb2", "LogProto"],
    "run_log": ["log_pb2", "LogProto"],
    "multiple_raw": ["multiple_raw_pb2", "MultipleRawProto"],
    "online_calib": ["online_calib_msg_pb2", "OnlineCalibMsgProto"],
    "crop_roi": ["perception_base_pb2", "Rect"],
    "parsing": ["parsing_pb2", "Parsing"],
    "lane_parsing": ["parsing_pb2", "Parsing"],
    "parsing_conf": ["extra_msg_info_pb2", "Image"],
    "depth_map": ["extra_msg_info_pb2", "DepthMap"],
    "raw": ["raw_pb2", "RawProto"],
    "roi_info": ["dts_pb2", "RoiInfo"],
    "runtime_env": ["runtime_env_pb2", "RuntimeEnvProto"],
    "runtime": ["runtime_pb2", "RuntimeProto"],
    "vehicledata": ["can_pb2", "VehicleMsg"],
    "can": ["can_pb2", "CANFrames"],
    "ethcan": ["usbcan_pb2", "EthernetFrame"],
    "eq4": ["usbcan_pb2", "EthernetFrame"],
    "novatel": ["gnss_pb2", "GNSSFrame"],
    "ub482": ["gnss_pb2", "GNSSFrame"],
    "ekv_f9k": ["gnss_pb2", "GNSSFrame"],
    "f9p": ["gnss_pb2", "GNSSFrame"],#TODO F9P
    "gnssodo": ["gnss_odo_pb2", "Odometry"],
    "xsens": ["imu_pb2", "IMUFrame"],
    "f9k_imu": ["imu_pb2", "IMUFrame"],
    "asm330": ["imu_pb2", "IMUFrame"],#todo F9Kimu
    "robosense": ["lidar_pb2", "LidarFrame"],
    "ibeo": ["lidar_pb2", "LidarFrame"],
    "hesai": ["lidar_pb2", "LidarFrame"],
    "velodyne": ["lidar_pb2", "LidarFrame"],
    "sick": ["lidar_pb2", "LidarFrame"],
    "continental": ["radar_pb2", "RadarFrame"],
    "usbcan": ["usbcan_pb2", "USBCANFrameRaw"],
    "odometry": ["odometry_pb2", "OdometryFrame"],
    "postinput_runmsg": ["odometry_pb2", "OdometryFrame"],
    "obstacles_raw": ["obstacle_raw_pb2", "ObstacleRaws"],
    "skeleton_raw": ["skeleton_pb2", "SkeletonRaws"],
    "camera_runtime": ["camera_pb2", "CameraParam"],
    "camera_default": ["camera_pb2", "CameraParam"],
    "fake_image_data": ["image_pb2", "Image"],
    "scan_points": ["scan_point_pb2", "ScanPoints"],
    "real3d": ["real3d_pb2", "Real3DProto"],
    "boxes_3d_raw": ["box_3D_raw_pb2", "Box3DRaws"],
    "gps_raw": ["sensor_pb2", "GPSFrameRaw"],
    "lidar_perception": ["obstacle_pb2", "Obstacles"],
    "lidar_perception_raw": ["obstacle_pb2", "Obstacles"],
    "adas_result": ["mono_adas_output_pb2", "ADASResultFrame"],
    "adas_hmi": ["mono_adas_hmi_pb2", "AdasHmi"],
    "sync_obstacle": ["obstacle_pb2", "Obstacles"],
    "obstacle": ["obstacle_pb2", "Obstacles"],
    "obs_track_run": ["obstacle_pb2", "ObstacleTrackRun"],
    "obs_state_run": ["obstacle_pb2", "ObstacleStateRun"],
    "obstacle_v2": ["obstacle", "v2_pb2", "Obstacles"],
    "data_type": ["obstacle_pb2", "Obstacles"],
    "bbox": ["obstacle_pb2", "Obstacles"],
    "traffic_sign": ["obstacle_pb2", "Obstacles"],
    "traffic_light": ["obstacle_pb2", "Obstacles"],
    "parsing_bbox": ["obstacle_pb2", "Obstacles"],
    "parsing_bbox_v2": ["obstacle_pb2", "Obstacles"],
    "ped_result": ["obstacle_pb2", "Obstacles"],
    "vehicle_result": ["obstacle_pb2", "Obstacles"],
    "speed_camera": ["obstacle_pb2", "Obstacles"],
    "sync_line": ["lane_pb2", "Lines"],
    "lane_result": ["lane_pb2", "Lines"],
    "pole_result": ["lane_pb2", "Lines"],
    "general_lane_result": ["lane_pb2", "Lines"],
    "road_height": ["lane_pb2", "Lines"],
    "resflow": ["parsing_pb2", "Parsing"],
    "elevation_gamma": ["parsing_pb2", "Parsing"],
    "bev_parsing": ["parsing_pb2", "Parsing"],
    "online_cali_status": ["parsing_pb2", "Parsing"],
    "loss_cali_status": ["parsing_pb2", "Parsing"],
    "laneline_v2": ["line", "v2_pb2", "Lines"],
    "cam_freespace": ["freespace", "v2_pb2", "CamContourPointsMsg"],
    "freespace_fusion": ["freespace", "v2_pb2", "FreespacePointsMsg"],
    "object_light": ["object_pb2", "Objects"],
    "object_sign": ["object_pb2", "Objects"],
    "object_parsing_bbox": ["object_pb2", "Objects"],
    # planning and control
    "control_cmd": ["control_pb2", "Control"],
    "vehicle_status": ["can_pb2", "VehicleStatus"],
    "radar_objects": ["obstacle_pb2", "Obstacles"],
    "radar_filter_objects": ["obstacle_pb2", "Obstacles"],
    "vision_objects": ["obstacle_pb2", "Obstacles"],
    "multisensor_fusion_objects": ["obstacle_pb2", "Obstacles"],
    "global_odometry": ["ad_base_pb2", "Odometry"],
    "vehicle_odometry": ["ad_base_pb2", "Odometry"],
    "ssm_model": ["ssm_pb2", "StandardScenarioModel"],
    "ssm_road": ["ssm_road", "v2_pb2", "Road"],
    "planning_traj": ["planning_pb2", "PlanningTraj"],
    "planning_path": ["planning_pb2", "PlanningTraj"],
    "chassis_can": ["can_pb2", "CANFrames"],
    "radar_can": ["can_pb2", "CANFrames"],
    "traffic_line_frame": ["ssm_pb2", "TrafficLinesFrame"],
    "ad_runtime_env": ["runtime_env_pb2", "RuntimeEnvProto"],
    "control_info": ["control_pb2", "ControlInfo"],
    # "ad_ndm_map_v2": ["ndm_pb2","NDMMessage"],//old
    "ad_ndm_map_v2": ["map_env","v2_pb2","MapEnvMsg"],
    "ad_ndm_nav": ["navigation","v2_pb2","NavigationMsg"],
    "ad_ndm_loc": ["location","v2_pb2","LocationMsg"],
    "lidar_traj": ["obstacle_pb2","Obstacles"],
    "function_state_to_hmi": ["state_manager_pb2","FunctionStateToHmiMessage"],
    "ad_state_manager": ["state_manager_pb2","StateManagerOutMessage"],
    "ad_noa_state": ["state_manager_pb2","NOAStateMessage"],
    "ad_adas_state": ["state_manager_pb2","ADASStateMessage"],
    "ssm_to_agsm": ["ssm_to_agsm_pb2","SSMFrameType"],

    # location and map
    "location": ["location","v2_pb2","LocationMsg"],
    "map_env_v2": ["map_env","v2_pb2","MapEnvMsg"],
    "rec_env_v2": ["map_env","v2_pb2","MapEnvMsg"],
    "fus_nav_v2": ["navigation","v2_pb2","NavigationMsg"],
    "fus_env_v2": ["map_env","v2_pb2","MapEnvMsg"],
    "navigation": ["navigation","v2_pb2","NavigationMsg"],
    "loc_debug_msg": ["loc_debug","v2_pb2","LocDebugMsg"],
    "loc_parsing": ["parsing_pb2", "Parsing"],
    "loc_lane_parsing": ["parsing_pb2", "Parsing"],
    # hmi
    "hmi_fsd_obstacle": ["obstacle_pb2", "Obstacles"],
    "hmi_fsd_points_bev": ["render_pb2", "DrawNode"],
    "hmi_fsd_points_pred": ["render_pb2", "DrawNode"],
    "hmi_pilot_lines": ["lane_pb2", "Lines"],
    "hmi_pilot_map": [], # NOT FOUND
    "hmi_pilot_location": [], # NOT FOUND
    "hmi_pilot_obstacle": ["obstacle_pb2", "Obstacles"],
    # generic render
    "render": ["render_pb2", "DrawNode"],
    #bev
    "bev3d": ["real3d_pb2", "Real3DProto"],
    "multi_cam_state": ["diag_info_pb2", "MultiCamState"],
    #radar
    "radar": ["radar","v2_pb2", "RadarFrame"],
    #HMI
    "hmi_drive": ["hmi", "hmi_drive_pb2", "DrivingMsg"]
}


class Chunk(object):
    def __init__(self, reader, offset, length, lazy):
        self._reader = reader
        self.offset = offset
        self.length = length
        self.content = None
        if not lazy:
            self.read()

    def read(self):
        if self.content is None:
            self.content = self._reader._read(self.offset, self.length)
        return self.content


class Data(Chunk):
    def __init__(self, reader, offset, length, data_descriptor, topic, lazy):
        super(Data, self).__init__(reader, offset, length, lazy)
        self.data_descriptor = data_descriptor
        self.topic = topic


class Frame(Chunk):
    def __init__(self, reader, meta, offset, length, data_list, topic_meta):
        super(Frame, self).__init__(reader, offset, length, lazy=True)
        self.meta = meta
        self.data_list = data_list
        self.topic_meta = topic_meta


class DataDescriptor(object):
    def __init__(self, type):
        self.type = type


class Reader(object):
    def __init__(self, fid, topics=[], lazy=True):
        self._fid = fid
        self.version = struct.unpack('i', self._fid.read(4))[0]
        self._lazy = lazy
        self.topics = topics
        self.topics_obj = {}
        for topic in topics:
            if topic in TOPIC_PROTOS:
                topic_proto = TOPIC_PROTOS[topic]
                module = importlib.import_module("pack_to_json.protocol.%s" % ".".join(topic_proto[:-1]))
                self.topics_obj[topic] = getattr(module, topic_proto[-1])

    def __del__(self):
        self.close()
        pass

    def close(self):
        self._fid.close()
        pass

    def _read(self, offset, length):
        old_pos = self._fid.tell()
        self._fid.seek(offset)
        buf = self._fid.read(length)
        self._fid.seek(old_pos)
        return buf

    def _parse_frame_v0(self, frame_meta):
        data_list = []
        ss = self._fid.read(4)
        if ss == '':
            status = END
        else:
            length = struct.unpack('i', ss)[0]
            offset = self._fid.tell()
            if length == 0:
                status = CONTINUE
            else:
                if frame_meta.frame_v0.img_frame.count:
                    image_count = frame_meta.frame_v0.img_frame.count
                else:
                    image_count = 1
                left_length = length
                for i in range(image_count):
                    ss = self._fid.read(4)
                    if ss == '':
                        status = END
                        break
                    d_offset = self._fid.tell()
                    d_length = struct.unpack('i', ss)[0]
                    data = Data(self, d_offset, d_length, DataDescriptor("image"), "", self._lazy)
                    data_list.append(data)
                    self._fid.seek(d_offset + d_length)
                    left_length -= (d_length + 4)
                if left_length > 0:
                    for i in range(image_count):
                        ss = self._fid.read(4)
                        if ss == '':
                            status = END
                            break
                        d_offset = self._fid.tell()
                        d_length = struct.unpack('i', ss)[0]
                        data = Data(self, d_offset, d_length, DataDescriptor("parsing"), "", self._lazy)
                        data_list.append(data)
                        self._fid.seek(d_offset + d_length)
                self._fid.seek(offset + length)
            status = SUCCESS
        return status, data_list

    def _parse_frame_v1(self, frame_meta):
        data_list = []
        ss = self._fid.read(4)
        if ss == '':
            status = END
        else:
            length = struct.unpack('i', ss)[0]
            offset = self._fid.tell()
            if length == 0:
                status = SUCCESS
            else:
                for data_descriptor in frame_meta.frame_v1.data_descriptor:
                    if hasattr(data_descriptor, "data") and data_descriptor.data.with_data_field == False:
                        continue
                    ss = self._fid.read(4)
                    if ss == '':
                        status = END
                        break
                    d_offset = self._fid.tell()
                    d_length = struct.unpack('i', ss)[0]
                    data = Data(self, d_offset, d_length, data_descriptor, "", self._lazy)
                    data_list.append(data)
                    self._fid.seek(d_offset + d_length)
                self._fid.seek(offset + length)
                status = SUCCESS
        return status, data_list

    def _parse_frame_v2(self, frame_meta):
        data_list = []
        topic_meta = {}
        ss = self._fid.read(4)
        if ss == '':
            status = END
        else:
            length = struct.unpack('i', ss)[0]
            offset = self._fid.tell()
            if length == 0:
                status = SUCCESS
                for data_descriptor in frame_meta.data_list.data_descriptor:
                    topic_name = data_descriptor.message_name
                    if not hasattr(data_descriptor, "proto") or len(data_descriptor.proto) == 0:
                        continue

                    if topic_name in self.topics:
                        topic_meta[topic_name] = []
                    for serialized_proto in data_descriptor.proto:
                        if hasattr(serialized_proto, "meta") and topic_name in topic_meta and topic_name in self.topics_obj:
                            cur_obj = self.topics_obj[topic_name]()
                            cur_obj.ParseFromString(serialized_proto.meta)
                            topic_meta[topic_name].append(cur_obj)
            else:
                for data_descriptor in frame_meta.data_list.data_descriptor:
                    topic_name = data_descriptor.message_name
                    if not hasattr(data_descriptor, "proto") or len(data_descriptor.proto) == 0:
                        continue

                    if topic_name in self.topics:
                        topic_meta[topic_name] = []
                    for serialized_proto in data_descriptor.proto:
                        if hasattr(serialized_proto, "meta") and topic_name in topic_meta and topic_name in self.topics_obj:
                            cur_obj = self.topics_obj[topic_name]()
                            cur_obj.ParseFromString(serialized_proto.meta)
                            topic_meta[topic_name].append(cur_obj)
                        if hasattr(serialized_proto, "data_index") and serialized_proto.data_index > -1:
                            ss = self._fid.read(4)
                            if ss == '':
                                status = END
                                break
                            d_offset = self._fid.tell()
                            d_length = struct.unpack('i', ss)[0]
                            data = Data(self, d_offset, d_length, serialized_proto, data_descriptor.message_name, self._lazy)
                            data_list.append(data)
                            self._fid.seek(d_offset + d_length)
                self._fid.seek(offset + length)
                status = SUCCESS
        return status, data_list, topic_meta

    def __next__(self):
        while True:
            # get_frame
            try:
                ss = self._fid.read(4)
                if ss == '':
                    raise StopIteration
                offset = self._fid.tell()
                length = struct.unpack('i', ss)[0]
                if length == 0:
                    raise StopIteration

                meta = self._fid.read(length)
            except Exception as e:
                raise StopIteration

            frame_meta = Frame_v2()
            try:
                frame_meta.ParseFromString(meta)
            except Exception as e:
                logging.warning('Parse frame failed')
                raise StopIteration
            version = frame_meta.version
            proto_version = frame_meta.proto_version
            topic_meta = {}
            if proto_version == 2:
                try:
                    status, data_list, topic_meta = self._parse_frame_v2(frame_meta)
                except Exception as e:
                    import traceback
                    traceback.print_exc()
                    logging.warning('Parse frame v2 failed')
                    raise StopIteration
            else:
                if version < 0:
                    frame_meta = FrameShuffle()
                else:
                    frame_meta = Frame_v1()
                try:
                    frame_meta.ParseFromString(meta)
                except Exception as e:
                    logging.warning('Parse frame failed')
                    raise StopIteration
                if proto_version == 0:
                    try:
                        status, data_list = self._parse_frame_v0(frame_meta)
                    except Exception as e:
                        logging.warning('Parse frame v0 failed')
                        raise StopIteration
                if proto_version == 1:
                    try:
                        status, data_list = self._parse_frame_v1(frame_meta)
                    except Exception as e:
                        logging.warning('Parse frame v1 failed')
                        raise StopIteration

            if status == END:
                raise StopIteration
            elif status == CONTINUE:
                continue
            elif status == SUCCESS:
                frame = Frame(self, frame_meta, offset, length, data_list, topic_meta)
                return frame
            else:
                logging.exception('Not implemented')
                raise NotImplementedError

    def reset(self):
        self._fid.seek(0)
        _ = struct.unpack('i', self._fid.read(4))[0]

    def __iter__(self):
        self.reset()
        return self
